package com.wisdom.system.api.core.aop;

import com.alibaba.fastjson.JSONObject;
import com.google.common.util.concurrent.RateLimiter;
import com.wisdom.system.entity.out.ResultBody;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * 并发限流切面
 *
 * @author ghh
 * 2019年12月24日 14:18:37
 */
@Scope
@Aspect
@Component
public class RateLimitAspect {

    @Autowired
    private HttpServletResponse response;

    //设置"并发数"为5,表示桶容量为 5 且每秒新增 5 个令牌，即每隔 200 毫秒新增一个令牌
    private RateLimiter rateLimiter = RateLimiter.create(5.0);

    @Pointcut("@annotation(com.wisdom.system.common.annotation.RateLimitAspect)")
    public void serviceLimit() {}

    @Around("serviceLimit()")
    public ResultBody around(ProceedingJoinPoint joinPoint) {
        //尝试获取一个令牌
        boolean flag = rateLimiter.tryAcquire();
        ResultBody resultBody = new ResultBody();
        try {
            if (flag) {
                resultBody = (ResultBody)joinPoint.proceed();
            }else{
                resultBody.setCode("-1");
                resultBody.setMsg("网络频繁！");
                String result = JSONObject.toJSONString(resultBody);
                output(response, result);
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return resultBody;
    }

    private void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes("UTF-8"));
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            outputStream.flush();
            outputStream.close();
        }
    }

}



